import os
import telnetlib

from datetime import datetime

import influxdb_client, os, time
from influxdb_client import InfluxDBClient, Point, WritePrecision
from influxdb_client.client.write_api import SYNCHRONOUS

# 设置
HOST = "192.168.2.98"
PORT = "23"
USER = b"root\n"
PASSWORD = b"vp202401\n"
COMMAND = b"./tmp/nvr_auto_script/performance_test/get_system_usage.sh\n"

def get_system_usage(host, port, user, password, command):
    try:
        with telnetlib.Telnet(host, port, timeout=10) as tn:
            tn.read_until(b"login: ")
            tn.write(user + b"\n")
            tn.read_until(b"Password: ")
            tn.write(password + b"\n")
            tn.write(command + b"\n")
            time.sleep(1)  # 等待命令执行和输出

            output = tn.read_very_eager().decode('ascii')

            cpu_usage = memory_usage = None
            for line in output.splitlines():
                if "CPU_USAGE=" in line:
                    cpu_usage = float(line.split("=")[1])
                elif "MEM_USAGE=" in line:
                    memory_usage = float(line.split("=")[1])

            return cpu_usage, memory_usage
    except Exception as e:
        print(f"Error: {e}")
        return None, None


def save_to_influxdb():
    token = 'I30C19snwIkjSFoolnHA2K47MwHr-i33DwuQr8aVaKst6ibj7YQzlhBeFcl-ac2KW8F9Z3jxlk5k7fXDubKzhQ=='
    org = "nvr_team"
    url = "http://localhost:8086"
    bucket = "nvr_bucket"
    # 初始化InfluxDB客户端
    client = InfluxDBClient(url=url, token=token, org=org)
    write_api = client.write_api(write_options=SYNCHRONOUS)
    # 创建数据点
    while True:
        cpu_usage, memory_usage = get_system_usage(HOST, PORT, USER, PASSWORD, COMMAND)
        # 创建数据点
        point_cpu = (
            Point("nvr_performance")
                .tag("resource", "cpu")
                .field("usage", cpu_usage)
                .time(datetime.utcnow())
        )

        point_memory = (
            Point("nvr_performance")
                .tag("resource", "memory")
                .field("usage", memory_usage)
                .time(datetime.utcnow())
        )

        # 写入数据到InfluxDB
        write_api.write(bucket=bucket, org=org, record=[point_cpu, point_memory])
        time.sleep(10)


if __name__ == "__main__":
    save_to_influxdb()
